//
// Created by hzj on 24-11-12.
//
#include "multi_camera_cooperation/preprocess.h"

Preprocess::Preprocess(ros::NodeHandle &nh)
{
    printf(BOLDREDPURPLE "[Preprocess] Start to preprocess...\n" RESET);

    // basic parameters
    nh.param<std::string>("cam", camera_name, "camA");
    nh.param<std::string>("servogroup", servogroup_name, "servogroup12");
    nh.param<std::string>("coopestimation_topic", coopestimation_topic, " ");

    // System Initialization
    sub_T_base_to_coopestimation = nh.subscribe(coopestimation_topic, 1, &Preprocess::T_base_to_coopestimation_callback, this);
    sub_T_base_to_servogroup = nh.subscribe("/T_base_to_" + servogroup_name, 1, &Preprocess::T_base_to_servogroup_callback, this);
    sub_T_servogroup_to_cam = nh.subscribe("/T_" + servogroup_name + "_to_" + camera_name, 1, &Preprocess::T_servogroup_to_cam_callback, this);
    sub_T_cam_to_estimation = nh.subscribe("/" + camera_name + "/single_cam_process_ros/ir_mono/T_cam_to_estimation", 1, &Preprocess::T_cam_to_estimation_callback, this);
    sub_OpticalFlowReadyFlag = nh.subscribe("/" + camera_name + "/single_cam_process_ros/ir_mono/opticalReadyFlag", 1, &Preprocess::OpticalFlowReadyFlag_callback, this);
    pub_T_base_to_cam_cal_raw = nh.advertise<geometry_msgs::TransformStamped>("/T_base_to_" + camera_name + "_cal_raw", 1);
    pub_T_base_to_cam = nh.advertise<geometry_msgs::TransformStamped>("/T_base_to_" + camera_name, 1);
    pub_T_base_to_cam_cal = nh.advertise<geometry_msgs::TransformStamped>("/T_base_to_" + camera_name + "_cal", 1);
    // pub_T_base_to_servogroup = nh.advertise<geometry_msgs::TransformStamped>("/T_base_to_" + servogroup_name, 1);
}

void Preprocess::T_base_to_servogroup_callback(const geometry_msgs::TransformStamped::ConstPtr& msg){
    auto q = msg->transform.rotation;
    auto t = msg->transform.translation;
    T_base_to_servogroup.block<3, 1>(0, 3) = Eigen::Vector3d(t.x, t.y, t.z);
    T_base_to_servogroup.block<3, 3>(0, 0) = Eigen::Quaterniond(q.w, q.x, q.y, q.z).toRotationMatrix();

    // std::cout << "T_base_to_servogroup: " << std::endl;
    // std::cout << T_base_to_servogroup << std::endl;
}

void Preprocess::T_servogroup_to_cam_callback(const geometry_msgs::TransformStamped::ConstPtr& msg){
    auto q = msg->transform.rotation;
    auto t = msg->transform.translation;
    T_servogroup_to_cam.block<3, 3>(0, 0) = Eigen::Quaterniond(q.w, q.x, q.y, q.z).toRotationMatrix();
    T_servogroup_to_cam.block<3, 1>(0, 3) = Eigen::Vector3d(t.x, t.y, t.z);

    // std::cout << "T_servogroup_to_cam: " << std::endl;
    // std::cout << T_servogroup_to_cam << std::endl;

    publish_T_base_to_cam();
    // publish_T_base_to_servogroup();
}

void Preprocess::T_base_to_coopestimation_callback(const geometry_msgs::PoseStamped::ConstPtr& msg){
    auto q = msg->pose.orientation;
    auto t = msg->pose.position;
    T_base_to_coopestimation.block<3, 3>(0, 0) = Eigen::Quaterniond(q.w, q.x, q.y, q.z).toRotationMatrix();
    T_base_to_coopestimation.block<3, 1>(0, 3) = Eigen::Vector3d(t.x, t.y, t.z);
}

void Preprocess::T_cam_to_estimation_callback(const geometry_msgs::TransformStamped::ConstPtr& msg){
    auto q = msg->transform.rotation;
    auto t = msg->transform.translation;
    T_cam_to_estimation.block<3, 3>(0, 0) = Eigen::Quaterniond(q.w, q.x, q.y, q.z).toRotationMatrix();
    T_cam_to_estimation.block<3, 1>(0, 3) = Eigen::Vector3d(t.x, t.y, t.z);
}

void Preprocess::publish_T_base_to_cam(){
    std::cout << std::endl;
    printf(BOLDREDPURPLE "[Preprocess] Start to generate T_base_to_cam...\n" RESET);

    if(T_servogroup_to_cam == Eigen::Matrix4d::Identity() || T_base_to_servogroup == Eigen::Matrix4d::Identity()){
        ROS_WARN("T_servogroup_to_cam or T_base_to_servogroup is not initialized!");
        return;
    }

    T_base_to_cam_pre = T_base_to_cam;

    T_base_to_cam = T_base_to_servogroup * T_servogroup_to_cam;

    tf::Vector3 t_base_to_cam = EigenVector3dToTFVector3(T_base_to_cam.block<3, 1>(0, 3));
    tf::Quaternion q_base_to_cam = EigenQuaterniondToTFQuaternion(Eigen::Quaterniond(T_base_to_cam.block<3, 3>(0, 0)));

    // std::cout << "T_base_to_cam: " << std::endl;
    // std::cout << T_base_to_cam << std::endl;

    T_base_to_cam_tf.setOrigin(t_base_to_cam);
    T_base_to_cam_tf.setRotation(q_base_to_cam);
    T_base_to_cam_tf.stamp_ = ros::Time::now();
    T_base_to_cam_tf.frame_id_ = "base";
    T_base_to_cam_tf.child_frame_id_ = camera_name;
    geometry_msgs::TransformStamped msg_T_base_to_cam;
    tf::transformStampedTFToMsg(T_base_to_cam_tf, msg_T_base_to_cam);
    pub_T_base_to_cam.publish(msg_T_base_to_cam);

    publish_T_base_to_cam_cal(OpticalFlowReadyFlag);
}

void Preprocess::publish_T_base_to_servogroup(){
    // temp
    T_base_to_cam = T_base_to_coopestimation;

    if (T_servogroup_to_cam == Eigen::Matrix4d::Identity() || T_base_to_cam == Eigen::Matrix4d::Identity())
    {
        ROS_WARN("T_base_to_coopestimation is not initialized!");
        return;
    }
    else
    {
        T_base_to_servogroup = T_base_to_cam * T_servogroup_to_cam.inverse();
    }

    tf::Vector3 t_base_to_servogroup = EigenVector3dToTFVector3(T_base_to_servogroup.block<3, 1>(0, 3));
    tf::Quaternion q_base_to_servogroup = EigenQuaterniondToTFQuaternion(Eigen::Quaterniond(T_base_to_servogroup.block<3, 3>(0, 0)));

    T_base_to_servogroup_tf.setOrigin(t_base_to_servogroup);
    T_base_to_servogroup_tf.setRotation(q_base_to_servogroup);
    T_base_to_servogroup_tf.stamp_ = ros::Time::now();
    T_base_to_servogroup_tf.frame_id_ = "base";
    T_base_to_servogroup_tf.child_frame_id_ = servogroup_name;
    geometry_msgs::TransformStamped msg_T_base_to_servogroup;
    tf::transformStampedTFToMsg(T_base_to_servogroup_tf, msg_T_base_to_servogroup);
    pub_T_base_to_servogroup.publish(msg_T_base_to_servogroup);

    if(true){
        if(count == 0){
            mean_t = t_base_to_servogroup;
            mean_q = q_base_to_servogroup;
            count++;
        }else{
            mean_t = (mean_t * count + t_base_to_servogroup) / (count + 1);
            mean_q = (mean_q * count + q_base_to_servogroup) / (count + 1);
            count++;
        }

        if(count < 2000){
            std::cout << "count: " << count << std::endl;
        }

        if(count == 2000){
            printf(GREEN "[Preprocess] mean_t: %f, %f, %f\n" RESET, mean_t.x(), mean_t.y(), mean_t.z());
            printf(GREEN "[Preprocess] mean_q: %f, %f, %f, %f\n" RESET, mean_q.x(), mean_q.y(), mean_q.z(), mean_q.w());
        }
    }
}

void Preprocess::publish_T_base_to_cam_cal(bool pnpGoodFlag){
    if(T_base_to_coopestimation == Eigen::Matrix4d::Identity())
    {
        ROS_WARN("T_base_to_coopestimation is not initialized!");
        return;
    }

    if(T_cam_to_estimation == Eigen::Matrix4d::Identity())
    {
        ROS_WARN("T_cam_to_estimation is not initialized!");
        return;
    }

    if(T_base_to_cam == Eigen::Matrix4d::Identity() || T_base_to_cam_pre == Eigen::Matrix4d::Identity())
    {
        ROS_WARN("T_base_to_cam is not initialized!");
        return;
    }

    Eigen::Matrix4d transform = T_base_to_cam_pre.inverse() * T_base_to_cam;
    T_base_to_cam_cal_raw = T_base_to_coopestimation * T_cam_to_estimation.inverse();

    T_base_to_cam_cal_raw_tf.setOrigin(EigenVector3dToTFVector3(T_base_to_cam_cal_raw.block<3, 1>(0, 3)));
    T_base_to_cam_cal_raw_tf.setRotation(EigenQuaterniondToTFQuaternion(Eigen::Quaterniond(T_base_to_cam_cal_raw.block<3, 3>(0, 0))));
    T_base_to_cam_cal_raw_tf.stamp_ = ros::Time::now();
    T_base_to_cam_cal_raw_tf.frame_id_ = "base";
    T_base_to_cam_cal_raw_tf.child_frame_id_ = camera_name + "_cal_raw";
    geometry_msgs::TransformStamped msg_T_base_to_cam_cal_raw;
    tf::transformStampedTFToMsg(T_base_to_cam_cal_raw_tf, msg_T_base_to_cam_cal_raw);
    pub_T_base_to_cam_cal_raw.publish(msg_T_base_to_cam_cal_raw);

    // std::cout << "T_base_to_cam_pre: " << std::endl;
    // std::cout << T_base_to_cam_pre << std::endl;
    // std::cout << "T_base_to_cam_cal_raw: " << std::endl;
    // std::cout << T_base_to_cam_cal_raw << std::endl;

    // Eigen::Vector3d t_base_to_cam_eigen = T_base_to_cam.block<3, 1>(0, 3);
    // Eigen::Vector3d t_base_to_cam_pre_eigen = T_base_to_cam_pre.block<3, 1>(0, 3);
    // Eigen::Vector3d t_base_to_cam_cal_raw_eigen = T_base_to_cam_cal_raw.block<3, 1>(0, 3);
    // Eigen::Quaterniond q_base_to_cam_eigen = Eigen::Quaterniond(T_base_to_cam.block<3, 3>(0, 0));
    // Eigen::Quaterniond q_base_to_cam_pre_eigen = Eigen::Quaterniond(T_base_to_cam_pre.block<3, 3>(0, 0));
    // Eigen::Quaterniond q_base_to_cam_cal_raw_eigen = Eigen::Quaterniond(T_base_to_cam_cal_raw.block<3, 3>(0, 0));

    // outfile.open("/home/hezijia/catkin_ws/src/multi_camera_cooperation/data/" + camera_name + "_T_base_to_servogroup.txt", ios::app);
    // outfile << t_base_to_cam_eigen.x() << " " << t_base_to_cam_eigen.y() << " " << t_base_to_cam_eigen.z() << " "
    //         << q_base_to_cam_eigen.x() << " " << q_base_to_cam_eigen.y() << " " << q_base_to_cam_eigen.z() << " " << q_base_to_cam_eigen.w() << " "
    //         << t_base_to_cam_pre_eigen.x() << " " << t_base_to_cam_pre_eigen.y() << " " << t_base_to_cam_pre_eigen.z() << " "
    //         << q_base_to_cam_pre_eigen.x() << " " << q_base_to_cam_pre_eigen.y() << " " << q_base_to_cam_pre_eigen.z() << " " << q_base_to_cam_pre_eigen.w() << " "
    //         << t_base_to_cam_cal_raw_eigen.x() << " " << t_base_to_cam_cal_raw_eigen.y() << " " << t_base_to_cam_cal_raw_eigen.z() << " "
    //         << q_base_to_cam_cal_raw_eigen.x() << " " << q_base_to_cam_cal_raw_eigen.y() << " " << q_base_to_cam_cal_raw_eigen.z() << " " << q_base_to_cam_cal_raw_eigen.w() << " "
    //         << std::endl;
    // outfile.close();

    Eigen::Vector3d t_base_to_cam_pre = T_base_to_cam_pre.block<3, 1>(0, 3);
    Eigen::Vector3d t_base_to_cam_cal_raw = T_base_to_cam_cal_raw.block<3, 1>(0, 3);
    Eigen::Quaterniond q_base_to_cam_pre = Eigen::Quaterniond(T_base_to_cam_pre.block<3, 3>(0, 0));
    Eigen::Quaterniond q_base_to_cam_cal_raw = Eigen::Quaterniond(T_base_to_cam_cal_raw.block<3, 3>(0, 0));

    double pre_weight = 0.7;
    double cal_raw_weight = 0.3 * pnpGoodFlag;
    t_base_to_cam_cal_raw = (pre_weight * t_base_to_cam_pre + cal_raw_weight * t_base_to_cam_cal_raw) / (pre_weight + cal_raw_weight);
    q_base_to_cam_cal_raw.x() = (pre_weight * q_base_to_cam_pre.x() + cal_raw_weight * q_base_to_cam_cal_raw.x()) / (pre_weight + cal_raw_weight);
    q_base_to_cam_cal_raw.y() = (pre_weight * q_base_to_cam_pre.y() + cal_raw_weight * q_base_to_cam_cal_raw.y()) / (pre_weight + cal_raw_weight);
    q_base_to_cam_cal_raw.z() = (pre_weight * q_base_to_cam_pre.z() + cal_raw_weight * q_base_to_cam_cal_raw.z()) / (pre_weight + cal_raw_weight);
    q_base_to_cam_cal_raw.w() = (pre_weight * q_base_to_cam_pre.w() + cal_raw_weight * q_base_to_cam_cal_raw.w()) / (pre_weight + cal_raw_weight);

    Eigen::Matrix4d T_base_to_cam_final = Eigen::Matrix4d::Identity();
    T_base_to_cam_final.block<3, 3>(0, 0) = q_base_to_cam_cal_raw.toRotationMatrix();
    T_base_to_cam_final.block<3, 1>(0, 3) = t_base_to_cam_cal_raw;
    T_base_to_cam_final = T_base_to_cam_final * transform;

    // std::cout << "T_base_to_cam_final: " << std::endl;
    // std::cout << T_base_to_cam_final << std::endl;

    tf::Vector3 t_base_to_cam_cal = EigenVector3dToTFVector3(T_base_to_cam_final.block<3, 1>(0, 3));
    tf::Quaternion q_base_to_cam_cal = EigenQuaterniondToTFQuaternion(Eigen::Quaterniond(T_base_to_cam_final.block<3, 3>(0, 0)));

    T_base_to_cam_tf.setOrigin(t_base_to_cam_cal);
    T_base_to_cam_tf.setRotation(q_base_to_cam_cal);
    T_base_to_cam_tf.stamp_ = ros::Time::now();
    T_base_to_cam_tf.frame_id_ = "base";
    T_base_to_cam_tf.child_frame_id_ = camera_name + "_cal";
    geometry_msgs::TransformStamped msg_T_base_to_cam_cal;
    tf::transformStampedTFToMsg(T_base_to_cam_tf, msg_T_base_to_cam_cal);
    pub_T_base_to_cam_cal.publish(msg_T_base_to_cam_cal);
}

void Preprocess::OpticalFlowReadyFlag_callback(const std_msgs::Bool::ConstPtr& msg){
    OpticalFlowReadyFlag = msg->data;
}